
import os
import re
import sys

root = os.path.dirname(os.path.dirname(__file__))
sys.path.append(root)
sys.path.append(os.path.dirname(root))

from _configs import hyp
from util_my import add_arg, replace_content
from project.yolov5.data import create_labels, refactor

def run(task='', data_dir='', cfg_id='', github_dir='', device='', weights='', out_dir='', logger=None,
    main_queue=None, sub_queue=None, split='', epochs=0, batch_size=0, workers='8', im_sz='', **kwargs):

    print("run1")
    names = kwargs.get('names', [])
    exist = kwargs.get('exist', True)
    augment = kwargs.get('augment', False)
    weights_dir = kwargs.get('weights_dir', '')

    scale = kwargs.get('scale', 0.5)
    hsv_s = kwargs.get('hsv_s', 0.7)
    hsv_v = kwargs.get('hsv_v', 0.4)
    hsv_h = kwargs.get('hsv_h', 0.015)
    fliplr = kwargs.get('fliplr', 0.5)
    mosaic = kwargs.get('mosaic', 1.0)
    translate = kwargs.get('translate', 0.1)

    git_dir, code_dir = refactor(task=task, data_dir=data_dir, github_dir=github_dir)
    create_labels(data_dir=data_dir, names=names, split=split)

    is_cls = task == 'classification' or 'cls' in data_dir
    is_seg = task == 'segmentation' or 'seg' in data_dir
    is_det = task == 'detection' or 'det' in data_dir

    device_count = len(device.split(','))
    weights = weights or os.path.join(weights_dir or git_dir, f'{cfg_id}.pt')
    weights = weights if 'yolov5' in cfg_id else cfg_id
    data_dir = data_dir if is_cls else os.path.join(data_dir, 'data.yaml')

    add_arg('--data', data_dir)
    add_arg('--device', device)
    add_arg('--model' if is_cls else '--weights', weights)
    add_arg('--epochs', epochs or None)
    add_arg('--project', out_dir or None)
    add_arg('--batch-size', batch_size or None)
    add_arg('--imgsz', im_sz or None)
    add_arg('--workers', workers)
    add_arg('--exist-ok' if exist else None)

    hyp_path = os.path.join(code_dir, 'data', 'hyps', 'hyp.scratch-low.yaml')
    if is_cls:
        name = 'classify'

        def replace_cls_train(content):
            content = re.sub("weights='IMAGENET1K_V1'", 'pretrained=pretrained', content)
            content = re.sub('(path=test_dir),', r"\1 if os.path.exists(test_dir) else data_dir / 'train',", content)
            content = re.sub('(\n( *)model = smart_DDP(model))', r'\2names = model.names\1\2model.names = names', content)

            if augment == False:
                content = re.sub("augment=True", 'augment=False', content)
                def replace_augmentations(content):
                    return re.sub('CenterCrop\(size\)', 'LetterBox(size)', content)
                replace_content(os.path.join(code_dir, 'utils', 'augmentations.py'), replace_augmentations)
            return content

        replace_content(os.path.join(code_dir, 'classify', 'train.py'), replace_cls_train)
    elif is_seg:
        name = 'segment'

        def replace_torch_utils(content):
            content = re.sub(', static_graph=True', ', find_unused_parameters=True', content)
            return content
        replace_content(os.path.join(code_dir, 'utils', 'torch_utils.py'), replace_torch_utils)
        
        def replace_hyp(content):
            content = content.replace('translate: 0.1', f'translate: {translate}')
            content = content.replace('fliplr: 0.5', f'fliplr: {fliplr}')
            content = content.replace('mosaic: 1.0', f'mosaic: {mosaic}')
            content = content.replace('hsv_h: 0.015', f'hsv_h: {hsv_h}')
            content = content.replace('hsv_s: 0.7', f'hsv_s: {hsv_s}')
            content = content.replace('hsv_v: 0.4', f'hsv_v: {hsv_v}')
            content = content.replace('scale: 0.5', f'scale: {scale}')
            return content
        replace_content(hyp_path, replace_hyp)

    elif is_det:
        name = ''

    train_script = os.path.join(code_dir, name, 'train.py')

    if device_count > 1:
        # TODO torch.distributed.launch to torchrun
        os.system(f" \
            CUDA_VISIBLE_DEVICES={device} \
            python -m torch.distributed.run \
            --nproc_per_node={device_count} \
            --master_port=$RANDOM {train_script} \
            {' '.join(sys.argv[1:])}"
        )
    else:
        while root in sys.path:
            sys.path.remove(root)

        print("run2")

        try:
            if is_cls:
                import downloads.yolov5.classify.train as train
            elif is_seg:
                import downloads.yolov5.segment.train as train
            elif is_det:
                import downloads.yolov5.train as train
        except Exception as e:
            print(str(e))
            sys.path.append(os.path.dirname(git_dir))
            sys.path.append(github_dir)
            print(git_dir, code_dir, github_dir)
            if is_cls:
                import yolov5.classify.train as train
            elif is_seg:
                import yolov5.segment.train as train
            elif is_det:
                import yolov5.train as train

        import yaml
        opt = train.parse_opt(known=True)
        opt.main = main_queue
        opt.sub = sub_queue
        try:
            opt.hyp = yaml.safe_load(open(hyp_path))
        except:
            opt.hyp = hyp

        try:
            train.main(opt)
        except Exception as e:
            if logger:
                logger.log(logger.ERROR, e, exception=True)
            if sub_queue:
                sub_queue.put({'err': str(e)})